package com.tous.cloud.gateway.bean;

import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ReactorLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @author mengwei
 * @description 自定义灰度负载均衡器 (GrayRoundRobinLoadBalancer)
 * @createDate 2025/7/21 17:37
 */
public class GrayRoundRobinLoadBalancer implements ReactorLoadBalancer<ServiceInstance> {

    private final String serviceId;

    private final ServiceInstanceListSupplier supplier;

    public GrayRoundRobinLoadBalancer(ServiceInstanceListSupplier supplier, String serviceId) {
        this.supplier = supplier;
        this.serviceId = serviceId;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        // 获取所有服务实例
        return supplier.get().next()
                .map(instances -> getInstanceResponse(instances, request));
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances, Request request) {
        // 1. 从请求中获取灰度标记
        boolean isGrayRequest = isGrayRequest(request);

        // 2. 过滤出符合条件的实例
        List<ServiceInstance> filteredInstances = instances.stream()
                .filter(instance -> isGrayInstance(instance) == isGrayRequest)
                .collect(Collectors.toList());

        // 3. 如果没有匹配实例，回退到普通实例
        if (filteredInstances.isEmpty()) {
            filteredInstances = instances.stream()
                    .filter(instance -> !isGrayInstance(instance))
                    .collect(Collectors.toList());
        }

        // 4. 简单的轮询选择
        if (!filteredInstances.isEmpty()) {
            int index = (int) (System.currentTimeMillis() % filteredInstances.size());
            return new DefaultResponse(filteredInstances.get(index));
        }

        return new EmptyResponse();
    }

    private boolean isGrayRequest(Request request) {
        if (request.getContext() instanceof RequestDataContext) {
            RequestDataContext context = (RequestDataContext) request.getContext();
            HttpHeaders headers = context.getClientRequest().getHeaders();
            return "true".equals(headers.getFirst("grap-tag"));
        }
        return false;
    }

    private boolean isGrayInstance(ServiceInstance instance) {
        return "true".equals(instance.getMetadata().get("grap-tag"));
    }
}
